package com.olink;

import java.io.*;
import java.net.Socket;
import java.net.SocketException;
import java.security.MessageDigest;
import java.util.Base64;
import java.util.Scanner;

/**
 * WebSocket客户端实现
 */
public class WebSocketClient {
    private Socket socket;
    private InputStream inputStream;
    private OutputStream outputStream;
    private String host;
    private int port;
    private boolean isConnected = false;

    public WebSocketClient(String host, int port) {
        this.host = host;
        this.port = port;
    }

    /**
     * 连接到WebSocket服务器
     */
    public void connect() throws IOException {
        // 创建TCP连接
        socket = new Socket(host, port);
        inputStream = socket.getInputStream();
        outputStream = socket.getOutputStream();

        // 生成随机的Sec-WebSocket-Key
        String secWebSocketKey = generateRandomKey();

        // 构建WebSocket握手请求
        String handshakeRequest = buildHandshakeRequest(secWebSocketKey);

        // 发送握手请求
        outputStream.write(handshakeRequest.getBytes("UTF-8"));
        outputStream.flush();

        // 接收并验证握手响应
        if (verifyHandshakeResponse(secWebSocketKey)) {
            System.out.println("WebSocket连接已建立");
            isConnected = true;

            // 启动接收消息的线程
            new Thread(this::receiveMessages).start();
        } else {
            System.out.println("WebSocket连接失败");
            close();
        }
    }

    /**
     * 发送消息到服务器
     */
    public void sendMessage(String message) throws IOException {
        if (!isConnected) {
            throw new IOException("未连接到服务器");
        }

        if(message.equals("quit")) {
            System.out.println("客户端正常退出");
            close();
            return;
        }

        byte[] payload = message.getBytes("UTF-8");
        ByteArrayOutputStream frame = new ByteArrayOutputStream();

        // 设置FIN和Opcode(文本帧)
        frame.write(0x81);

        // 设置负载长度
        if (payload.length <= 125) {
            frame.write(payload.length);
        } else if (payload.length <= 65535) {
            frame.write(126);
            frame.write((payload.length >> 8) & 0xFF);
            frame.write(payload.length & 0xFF);
        } else {
            frame.write(127);
            for (int i = 7; i >= 0; i--) {
                frame.write((payload.length >> (8 * i)) & 0xFF);
            }
        }

        // 添加掩码键(客户端必须掩码)
        byte[] maskKey = new byte[4];
        for (int i = 0; i < 4; i++) {
            maskKey[i] = (byte) (Math.random() * 256);
        }
        frame.write(maskKey);

        // 应用掩码并写入数据
        byte[] maskedPayload = new byte[payload.length];
        for (int i = 0; i < payload.length; i++) {
            maskedPayload[i] = (byte) (payload[i] ^ maskKey[i % 4]);
        }
        frame.write(maskedPayload);

        // 发送帧
        try {
            outputStream.write(frame.toByteArray());
            outputStream.flush();
        }catch (SocketException e){
            //输入quit时，抛出异常，这里捕捉异常然后关闭连接
            throw new SocketException("SocketException");
        }

    }

    /**
     * 关闭连接
     */
    public void close() {
        try {
            isConnected = false;
            if (socket != null && !socket.isClosed()) {
                socket.close();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 接收并处理消息
     */
    private void receiveMessages() {
        try {
            while (isConnected) {
                int b1 = inputStream.read();
                if (b1 == -1) {
                    System.out.println("服务器关闭了连接");
                    break;
                }

                int b2 = inputStream.read();
                boolean fin = (b1 & 0x80) != 0;
                int opcode = b1 & 0x0F;
                boolean masked = (b2 & 0x80) != 0;
                long payloadLength = b2 & 0x7F;

                // 读取扩展长度
                if (payloadLength == 126) {
                    payloadLength = ((long) inputStream.read() << 8) | inputStream.read();
                } else if (payloadLength == 127) {
                    payloadLength = 0;
                    for (int i = 0; i < 8; i++) {
                        payloadLength = (payloadLength << 8) | inputStream.read();
                    }
                }

                // 服务器消息不需要掩码
                if (masked) {
                    System.out.println("错误：服务器消息被掩码");
                    break;
                }

                // 读取消息内容
                byte[] payload = new byte[(int) payloadLength];

                int bytesRead = 0;
                while (bytesRead < payloadLength) {
                    int count = inputStream.read(payload, bytesRead, (int) (payloadLength - bytesRead));
                    if (count == -1) {
                        break;
                    }
                    bytesRead += count;
                }

                // 处理控制帧
                if (opcode == 0x08) { // 关闭帧
                    System.out.println("收到服务器的关闭帧");
                    close();
                    break;
                } else if (opcode == 0x09) { // Ping帧
                    System.out.println("收到Ping帧，发送Pong响应");
                    sendPong(payload);
                } else if (opcode == 0x01) { // 文本帧
                    String message = new String(payload, "UTF-8");

                    if (message.contains("客户端/"+this.host+":"+this.socket.getLocalPort()))
                        continue;
                    System.out.println("收到消息: " + message);
                }
            }
        } catch (IOException e) {
            if (isConnected) {
                System.out.println("接收消息时发生错误: " + e.getMessage());
            }else{

            }
        }
    }

    /**
     * 发送Pong响应
     */
    private void sendPong(byte[] payload) throws IOException {
        ByteArrayOutputStream frame = new ByteArrayOutputStream();
        frame.write(0x8A); // Pong帧

        if (payload.length <= 125) {
            frame.write(payload.length);
        } else {
            frame.write(126);
            frame.write((payload.length >> 8) & 0xFF);
            frame.write(payload.length & 0xFF);
        }

        frame.write(payload);
        outputStream.write(frame.toByteArray());
        outputStream.flush();
    }

    /**
     * 构建WebSocket握手请求
     */
    private String buildHandshakeRequest(String secWebSocketKey) {
        return "GET / HTTP/1.1\r\n" +
                "Host: " + host + ":" + port + "\r\n" +
                "Upgrade: websocket\r\n" +
                "Connection: Upgrade\r\n" +
                "Sec-WebSocket-Key: " + secWebSocketKey + "\r\n" +
                "Sec-WebSocket-Version: 13\r\n\r\n";
    }

    /**
     * 验证WebSocket握手响应
     */
    private boolean verifyHandshakeResponse(String secWebSocketKey) throws IOException {
        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
        String line;
        boolean upgradeFound = false;
        boolean connectionFound = false;
        String acceptHeader = null;

        while ((line = reader.readLine()) != null && !line.isEmpty()) {
            if (line.startsWith("Upgrade: ")) {
                upgradeFound = line.contains("websocket");
            } else if (line.startsWith("Connection: ")) {
                connectionFound = line.contains("Upgrade");
            } else if (line.startsWith("Sec-WebSocket-Accept: ")) {
                acceptHeader = line.substring("Sec-WebSocket-Accept: ".length()).trim();
            }
        }

        // 计算预期的Sec-WebSocket-Accept值
        String expectedAccept = generateExpectedAccept(secWebSocketKey);

        return upgradeFound && connectionFound && expectedAccept.equals(acceptHeader);
    }

    /**
     * 生成随机的Sec-WebSocket-Key
     */
    private String generateRandomKey() {
        byte[] randomBytes = new byte[16];
        for (int i = 0; i < 16; i++) {
            randomBytes[i] = (byte) (Math.random() * 256);
        }
        return Base64.getEncoder().encodeToString(randomBytes);
    }

    /**
     * 生成预期的Sec-WebSocket-Accept值
     */
    private String generateExpectedAccept(String secWebSocketKey) {
        try {
            String concatenated = secWebSocketKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
            MessageDigest digest = MessageDigest.getInstance("SHA-1");
            byte[] hash = digest.digest(concatenated.getBytes("UTF-8"));
            return Base64.getEncoder().encodeToString(hash);
        } catch (Exception e) {
            throw new RuntimeException("生成Sec-WebSocket-Accept失败", e);
        }
    }

    public static void main(String[] args) {
        WebSocketClient client = new WebSocketClient("127.0.0.1", 8080);

        try {
            client.connect();

            // 从控制台读取输入并发送消息
            Scanner scanner = new Scanner(System.in);
            while (true) {
                String message = scanner.nextLine();
                if ("exit".equalsIgnoreCase(message)) {
                    break;
                }
                client.sendMessage(message);
                if(client.isConnected){
                    System.out.println("发送消息: " + message);
                }else{
                    throw new IOException();
                }
            }

        } catch (IOException e){
            //客户端退出时，仍尝试发送消息，然后报错，然后走到这里
        }
        finally {
            if(client.isConnected){
                client.close();
            }
        }
    }
}